#!/usr/bin/env python3

import argparse
import re
import subprocess
import os
import platform
import sys

# Just test these boot to start with.  In the future we should at least
# do ostree upgrades with uefi etc.  But we don't really need the *full*
# suite...if podman somehow broke with nvme or uefi I'd be amazed and impressed.
BASIC_SCENARIOS = ["nvme=true", "firmware=uefi"]
BASIC_SCENARIOS_SECURE_BOOT = ["firmware=uefi-secure"]
arch = platform.machine()

cosa_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, cosa_dir)

from cosalib import cmdlib

basearch = cmdlib.get_basearch()

# Parse args and dispatch
parser = argparse.ArgumentParser()
parser.add_argument("--build", help="Build ID")
parser.add_argument("--basic-qemu-scenarios", help="Run the basic test across uefi-secure,nvme etc.", action='store_true')
parser.add_argument("--output-dir", help="Output directory")
parser.add_argument("--skip-secure-boot", help="Use with '--basic-qemu-scenarios' to skip the Secure Boot tests", action='store_true')
parser.add_argument("--upgrades", help="Run upgrade tests", action='store_true')
parser.add_argument("subargs", help="Remaining arguments for kola", nargs='*',
                    default=[])
args, unknown_args = parser.parse_known_args()

default_output_dir_base = os.environ.get('ARTIFACT_DIR', 'tmp')

if args.upgrades:
    default_cmd = 'run-upgrade'
    default_output_dir = os.path.join(default_output_dir_base, "kola-upgrade")
else:
    default_cmd = 'run'
    default_output_dir = os.path.join(default_output_dir_base, "kola")

# XXX: teach to kola to auto-detect based on prefix; see discussions in
# https://github.com/coreos/coreos-assembler/pull/85
kolaargs = ['kola']

r = re.compile("-p(=.+)?|--platform(=.+)?")
platformargs = list(filter(r.match, unknown_args))

if os.getuid() != 0 and len(platformargs) == 0:
    kolaargs.extend(['-p', 'qemu'])

if args.build is not None:
    kolaargs.extend(['--build', args.build])
outputdir = args.output_dir or default_output_dir
subargs = args.subargs or [default_cmd]
kolaargs.extend(subargs)
kolaargs.extend(unknown_args)

if args.basic_qemu_scenarios:
    if arch == "x86_64":
        os.mkdir(outputdir)  # Create the toplevel output dir
        for scenario in BASIC_SCENARIOS:
            kolaargs.extend(['--output-dir',
                             os.path.join(outputdir, scenario.replace('=', '-'))])
            subargs = kolaargs + ['--qemu-' + scenario, 'basic']
            print(subprocess.list2cmdline(subargs), flush=True)
            subprocess.check_call(subargs)
        if not args.skip_secure_boot:
            for scenario in BASIC_SCENARIOS_SECURE_BOOT:
                kolaargs.extend(['--output-dir',
                                 os.path.join(outputdir, scenario.replace('=', '-'))])
                # See https://issues.redhat.com/browse/COS-2000 - there's
                # some bug with shim/grub2 that fails with secure boot on < ~1300MiB of RAM.
                # But we're not going to block on that; real world OCP worker nodes are at least 16GiB etc.
                subargs = kolaargs + ['--qemu-' + scenario, 'basic'] + ["--qemu-memory", "1536"]
                print(subprocess.list2cmdline(subargs), flush=True)
                subprocess.check_call(subargs)
    else:
        # Basic qemu scenarios using nvme and uefi
        # are not supported on multi-arch
        kolaargs.extend(['--output-dir', outputdir])
        subargs = kolaargs + ['basic']
        print(subprocess.list2cmdline(subargs), flush=True)
        subprocess.check_call(subargs)
elif args.upgrades:
    kolaargs.extend(['--output-dir', outputdir])
    if '--qemu-image-dir' not in unknown_args:
        os.makedirs('tmp/kola-qemu-cache', exist_ok=True)
        kolaargs.extend(['--qemu-image-dir', 'tmp/kola-qemu-cache'])
    kolaargs.extend(['-v', '--find-parent-image'])
    print(subprocess.list2cmdline(kolaargs), flush=True)
    os.execvp('kola', kolaargs)
else:
    kolaargs.extend(['--output-dir', outputdir])
    # flush before exec; see https://docs.python.org/3.7/library/os.html#os.execvpe
    print(subprocess.list2cmdline(kolaargs), flush=True)
    os.execvp('kola', kolaargs)
